import os
import shutil
def preparinbdata(main_xml_file, main_img_file, train_size, val_size):
    for i in range(0, train_size):
        source_xml = main_xml_file + "/" + metarial[i] + ".xml"
        source_img = main_img_file + "/" + metarial[i] + ".png"
        
        mstring = metarial[i]
        train_destination_xml = "./dataset/mask/train/labels" + "/" + metarial[i] + ".xml"
        train_destination_png = "./dataset/mask/train/images" + "/" + metarial[i] + ".png"
        
        shutil.copy(source_xml, train_destination_xml)
        shutil.copy(source_img, train_destination_png)
    for n in range(train_size  , train_size + val_size):

        source_xml = main_xml_file + "/" + metarial[n] + ".xml"
        source_img = main_img_file + "/" + metarial[n] + ".png"

        mstring = metarial[n]
        val_destination_xml = "./dataset/mask/val/labels" + "/" + metarial[n] + ".xml"
        val_destination_png = "./dataset/mask/val/images" + "/" + metarial[n] + ".png"

        shutil.copy(source_xml, val_destination_xml)
        shutil.copy(source_img, val_destination_png)
        
if __name__ == '__main__':
    metarial = []
    for i in os.listdir("./dataset/images"):
        str = i[:-4]
        metarial.append(str)
    train_size = int(len(metarial) * 0.7)
    val_size = int(len(metarial) * 0.3)
    print("Sum of image: ", len(metarial))
    print("Sum of the train size: ", train_size)
    print("Sum of the val size: ", val_size)
    if not os.path.exists("./dataset/mask"):
        os.mkdir('./dataset/mask')
        os.mkdir('./dataset/mask/train')
        os.mkdir('./dataset/mask/val')
        os.mkdir('./dataset/mask/train/images')
        os.mkdir('./dataset/mask/train/labels')
        os.mkdir('./dataset/mask/val/images')
        os.mkdir('./dataset/mask/val/labels')
    preparinbdata(main_xml_file = "./dataset/annotations", 
              main_img_file = "./dataset/images",
              train_size = train_size, 
              val_size = val_size)
